#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef size_t ea_t;
#ifdef _WINDEF_
typedef  unsigned char uint8_t;
typedef  char int8_t;
typedef  unsigned short uint16_t;
typedef  short int16_t;
typedef  unsigned __int32 uint32_t;
typedef  __int32 int32_t;
typedef  unsigned __int64 uint64_t;
typedef  __int64 int64_t;


#ifdef _WIN32
//32
#elif _WIN64
//64
#endif

#define strcasecmp stricmp
#define strncasecmp strnicmp

#else
#include <unistd.h>
#include <time.h>
#include <sys/wait.h>
#include <sys/mman.h>
#include <dlfcn.h>
#include <spawn.h>
#ifdef __APPLE__
//macosx or ios
char shellcode[]="\xeb\x30\x48\x31\xc0\x48\x31\xff\x48\x31\xd2\x48\xff\xc0\x48\xc1\xe0\x19\x48\x83\xc0\x04\x48\xff\xc7\x5e\x48\x83\xc2\x0e\x0f\x05\x48\x31\xc0\x48\xff\xc0\x48\xc1\xe0\x19\x48\xff\xc0\x48\x31\xff\x0f\x05\xe8\xcb\xff\xff\xff\x48\x65\x6c\x6c\x6f\x2c\x20\x77\x6f\x72\x6c\x64\x21\x0a";
#define MAP_ANONYMOUS MAP_ANON
#define _POSIX_SPAWN_DISABLE_ASLR 0x0100
#define	_POSIX_SPAWN_ALLOW_DATA_EXEC 0x2000
#else
//linux
#include <sys/personality.h>
#if __WORDSIZE==64
//64 gnu
char shellcode[]="\xeb\x24\x48\x31\xc0\x48\x31\xff\x48\x31\xd2\x48\x83\xc0\x01\x48\x83\xc7\x01\x5e\x48\x83\xc2\x0e\x0f\x05\x48\x31\xc0\x48\x83\xc0\x3c\x48\x31\xff\x0f\x05\xe8\xd7\xff\xff\xff\x48\x65\x6c\x6c\x6f\x2c\x20\x77\x6f\x72\x6c\x64\x21\x0a";
#elif __WORDSIZE==32
//32 gnu
char shellcode[]="\x31\xc0\x83\xc8\x02\xcd\x80\x85\xc0\x75\x01\xc3\x31\xc0\x50\x68\x6e\x2f\x73\x68\x68\x2f\x2f\x62\x69\x89\xe3\x50\x53\x89\xe1\x89\xc2\x83\xc8\x0b\xcd\x80";
#endif
#endif

#endif

//char shellcode[] = "\xeb\xfe";
void test_stack_dep()
{
    void (*f)();
    char x[0x100]={0};
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    printf("pid:%d stack addr:%p\n",getpid(),&x);
    getchar();
    f();
    printf("stack dep end\n");

}

void test_stack_dep_close()
{
    void (*f)();
    char x[0x100]={0};
    memcpy(x, shellcode, sizeof(shellcode));
    int pagesize = getpagesize();
    ea_t base = (ea_t)x & ~(pagesize - 1);
    int ret=mprotect((void *) base, pagesize, PROT_READ | PROT_WRITE | PROT_EXEC);
    if(ret<0)
    { 
        perror("mprotect fail!\n");
        return;
    }
    f = (void (*)()) x;
    f();
    printf("stack dep close success!\n");

}

void test_heap_dep()
{
    void (*f)();
    char *x =(char*) malloc(sizeof(shellcode));
    printf("%p\n",x);
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    f();
}

void test_heap_dep_close()
{
    void (*f)();
    char *x =(char*) malloc(sizeof(shellcode));
    int pagesize = getpagesize();
    ea_t base = (ea_t)x & ~(pagesize - 1);
    int ret = mprotect((void *) base, pagesize, PROT_READ | PROT_WRITE | PROT_EXEC);
    if(ret<0)
    { 
        perror("mprotect fail!\n");
        return;
    }
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    f();
    printf("heap dep close success!\n");

}

void test_map_zero_addr()
{
    void (*f)();
    void *x=NULL;
    x=mmap(NULL,0x1000,PROT_READ|PROT_WRITE|PROT_EXEC,MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE,0,0);
    if(x==MAP_FAILED)
    {
        perror("mmap 0x0 addr fail!\n");
        x=mmap(NULL,0x1000,PROT_READ|PROT_WRITE|PROT_EXEC,MAP_ANONYMOUS|MAP_PRIVATE,0,0);
        if(x==MAP_FAILED)
        {
            perror("mmap any addr fail!fuck\n");
            return;
        }
    }
    printf("mmap other addr success:%p\n",x);
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    f(); 

    return;
}

void * get_process_baseaddr()
{
    Dl_info info;
    int ret=dladdr((void*)get_process_baseaddr,&info);
    if(ret==0)
    {
        printf("dladdr fail!\n");
        return NULL;
    }
    return info.dli_fbase;
}


void test_aslr()
{
    //printf("pid %d\n",getpid());
    int size=0x1000000;
    void * h=malloc(size);
    free(h);
    printf("heap:%p\n",h);
    printf("stack:%p\n",&size);
    void * baseaddr=get_process_baseaddr();
    printf("main:%p\n",baseaddr);
#ifdef __APPLE__
    void * handle=dlopen("libc.dylib",RTLD_LAZY);
#else
    void * handle=dlopen("libc.so.6",RTLD_LAZY);
#endif
    printf("lib:%p\n",handle);
    dlclose(handle);
    
    return;
}

void test_shellcode()
{
    void (*f)();
    void *x=mmap(NULL,0x1000,PROT_READ|PROT_WRITE|PROT_EXEC,MAP_ANONYMOUS|MAP_PRIVATE,0,0);
    if(x==MAP_FAILED)
    {
        perror("mmap any addr fail!fuck\n");
        return;
    }
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    f(); 
    return;
    /*
    int status=0;
    int options=0;
    int orig_persona=personality(-1);
    if(-1==orig_persona)
    {
        perror("personality fail!\n");
        return;
    }
    int ret=personality(orig_persona|READ_IMPLIES_EXEC);
    if(-1==ret)
    {
        perror("personality fail!\n");
        return;
    }
    void (*f)();
    char *x =(char*) malloc(sizeof(shellcode));
    memcpy(x, shellcode, sizeof(shellcode));
    f = (void (*)()) x;
    f(); 
    */
}

void usage(char * name)
{
    printf("%s -ts (test shellcode)\n",name);
    printf("%s -tm (test map zero addr and exec shellcode)\n",name);
    printf("%s -tss (test exec shellcode in stack)\n",name);
    printf("%s -tsh (test exec shellcode in heap)\n",name);
    printf("%s -tmss (test mprotect stack and exec shellcode in stack)\n",name);
    printf("%s -tmsh (test mprotect heap and exec shellcode in heap)\n",name);
    printf("%s -ta (test aslr)\n",name);
    return;
}

int main(int argc, char *argv[],char * env[])
{

    if(argc<2)
    {
        usage(argv[0]);
        return -1;
    }
    char * option=argv[1];
    
    if(strcasecmp(option,"-ts")==0)
        test_shellcode();
    else if(strcasecmp(option,"-tm")==0)
        test_map_zero_addr();
    else if(strcasecmp(option,"-tss")==0)
        test_stack_dep();
    else if(strcasecmp(option,"-tsh")==0)
        test_heap_dep();
    else if(strcasecmp(option,"-tmss")==0)
        test_stack_dep_close();
    else if(strcasecmp(option,"-tmsh")==0)
        test_heap_dep_close();
    else if(strcasecmp(option,"-ta")==0)
        test_aslr();
    else
        usage(argv[0]);

    return 0;

}
